package com.noahsloan.nutils.streams;

import java.io.FilterInputStream;
import java.io.IOException;
import java.io.InputStream;

/**
 * PartialInputStream wraps an InputStream, allowing only a certain number of
 * bytes to be read from it.
 * <p>
 * A PartialInputStream is even safe if the InputStream wrapped is null, as long
 * as the given limit is 0.
 * 
 * @author Noah Sloan
 * 
 */
public class PartialInputStream extends FilterInputStream {

    private long limit; // the number of bytes we are allowed to read

    private long read; // the number of bytes read

    /**
     * 
     * @param in
     *            the stream to read from.
     * @param limit
     *            the maximum number of bytes that may be read from in.
     */
    public PartialInputStream(InputStream in, long limit) {
        super(in);
        this.limit = limit;
        this.read = 0;
    }

    @Override
    /**
     * Operates on the general contract of read(). -1 indicates end of stream.
     */
    public int read() throws IOException {
        if (read < limit) {
            read++;
            return super.read();
        } else {
            return -1;
        }
    }

    @Override
    /**
     * @see InputStream#read(byte[])
     */
    public int read(byte[] buffer) throws IOException {
        return read(buffer, 0, buffer.length);
    }

    @Override
    /**
     * @see InputStream#read(byte[],int,int)
     */
    public int read(byte[] buffer, int start, int length) throws IOException {
        if (length > limit - read) {
            length = (int) (limit - read);
        }
        if (length > 0) {
            int i = super.read(buffer, start, length);
            read += i;
            return i;
        }
        return -1;
    }

    @Override
    /**
     * @see InputStream#available()
     * 
     * The difference between this and the superclass method is that if
     * super.available() returns more bytes than what is left of the limit, then
     * what is left of the limit is returned.
     */
    public int available() throws IOException {
        long foo = limit - read;
        return (int) Math.min(foo, super.available());
    }

    @Override
    /**
     * Consumes any remaining bytes up to the limit. Does not close the
     * underlying stream unless the end of the stream is reached.
     */
    public void close() throws IOException {
        while (read < limit) {
            if (read() == -1) {
                super.close();
                return;
            }
        }
    }

    /**
     * Returns an unbounded stream if limit is less than 0, otherwise returns a
     * {@link PartialInputStream} with the given limit.
     * 
     * @param in
     * @param limit
     * @return
     */
    public static InputStream getStream(InputStream in, long limit) {
        return limit < 0 ? in : new PartialInputStream(in, limit);
    }
}
